from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
import os
import random
import sys

import tensorflow as tf

from datasets import dataset_utils

_NUM_SHARDS = 4
_NUM_VALIDATION = 1000

class ImageReader(object):
    """Helper class that provides TensorFlow image coding utilities."""

    def __init__(self):
        # Initializes function that decodes RGB JPEG data.
        self._decode_jpeg_data = tf.placeholder(dtype=tf.string)
        self._decode_jpeg = tf.image.decode_jpeg(self._decode_jpeg_data, channels=3)
        self._decode_jpeg = tf.image.resize_images(self._decode_jpeg, (224, 224))

    def read_image_dims(self, sess, image_data):
        #image = self.decode_jpeg(sess, image_data)
        return 224, 224

    def decode_jpeg(self, sess, image_data):
        image = sess.run(self._decode_jpeg,
                         feed_dict={self._decode_jpeg_data: image_data})
        assert len(image.shape) == 3
        assert image.shape[2] == 3
        return image

def _get_filenames_and_classes(dataset_dir):
    text_file_path = os.path.join(dataset_dir, 'Flickr8k_text/Flickr8k.lemma.token.txt')
    text_content = open(text_file_path, 'r', encoding='utf-8').read()

    class_names = []
    photo_filenames = []
    row_data = text_content.split('\n')
    for row in row_data:
        arr = row.split('\t')
        if len(arr) > 1:
            image = arr[0].split('#')[0]
            captions = arr[1]
            photo_filenames.append(dataset_dir + 'Flickr8k_Dataset/' + image)
            class_names.append(captions)
    return photo_filenames, class_names

def _get_dataset_filename(dataset_dir, split_name, shard_id):
    output_filename = 'flickr_%s_%05d-of-%05d.tfrecord' % (
        split_name, shard_id, _NUM_SHARDS)
    return os.path.join(dataset_dir, output_filename)

def _convert_dataset(split_name, filenames, class_names, dataset_dir):
    assert split_name in ['train', 'validation']

    num_per_shard = int(math.ceil(len(filenames) / float(_NUM_SHARDS)))

    with tf.Graph().as_default():
        image_reader = ImageReader()

        with tf.Session('') as sess:

            for shard_id in range(_NUM_SHARDS):
                output_filename = _get_dataset_filename(
                    dataset_dir, split_name, shard_id)

                with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
                    start_ndx = shard_id * num_per_shard
                    end_ndx = min((shard_id + 1) * num_per_shard, len(filenames))
                    for i in range(start_ndx, end_ndx):
                        sys.stdout.write('\r>> Converting image %d/%d shard %d' % (
                            i + 1, len(filenames), shard_id))
                        sys.stdout.flush()

                        # Read the filename:
                        image_data = tf.gfile.FastGFile(filenames[i], 'rb').read()
                        height, width = image_reader.read_image_dims(sess, image_data)

                        class_name = class_names[i]

                        example = _image_to_tfexample(
                            image_data, b'jpg', height, width, class_name)
                        tfrecord_writer.write(example.SerializeToString())

    sys.stdout.write('\n')
    sys.stdout.flush()


def _image_to_tfexample(image_data, image_format, height, width, class_name):
  return tf.train.Example(features=tf.train.Features(feature={
      'image/encoded': dataset_utils.bytes_feature(image_data),
      'image/format': dataset_utils.bytes_feature(image_format),
      'image/class/label': dataset_utils.bytes_feature(class_name),
      'image/height': dataset_utils.int64_feature(height),
      'image/width': dataset_utils.int64_feature(width),
  }))

def _dataset_exists(dataset_dir):
    for split_name in ['train', 'validation']:
        for shard_id in range(_NUM_SHARDS):
            output_filename = _get_dataset_filename(
                dataset_dir, split_name, shard_id)
            if not tf.gfile.Exists(output_filename):
                return False
    return True

def run(dataset_dir):

    if not tf.gfile.Exists(dataset_dir):
        tf.gfile.MakeDirs(dataset_dir)

    if _dataset_exists(dataset_dir):
        print('Dataset files already exist. Exiting without re-creating them.')
        return

    photo_filenames, class_names = _get_filenames_and_classes(dataset_dir)

    training_filenames = photo_filenames[_NUM_VALIDATION:]
    validation_filenames = photo_filenames[:_NUM_VALIDATION]

    training_classnames = class_names[_NUM_VALIDATION:]
    validation_classnames = class_names[:_NUM_VALIDATION]

    # First, convert the training and validation sets.
    _convert_dataset('train', training_filenames, training_classnames,
                     dataset_dir)
    _convert_dataset('validation', validation_filenames, validation_classnames,
                     dataset_dir)

    print('\nFinished converting the Flowers dataset!')